import paddlehub as hub
from openpyxl import load_workbook

if __name__ == "__main__":
    # 加载senta模型
    senta = hub.Module(name="senta_bilstm")
    comments = list()
    with open("./fqly_comment.txt", "r", encoding='utf-8') as f:
        for line in f.readlines():
            line = line.strip("\n")
            if line != "":
                comments.append(line)
    # 把要测试的短文本以str格式放到这个列表里
    test_text = comments
    # 指定模型输入
    input_dict = {"text": test_text}
    # 把数据喂给senta模型的文本分类函数
    results = senta.sentiment_classify(data=input_dict)
    # 遍历分析每个短文本
    positive_nums = 0
    negative_nums = 0
    for index, text in enumerate(test_text):
        results[index]["text"] = text
    file_name = "./comment_emotion.xlsx"
    wb = load_workbook(filename=file_name)
    ws = wb.active
    fields = ["评论", "评论情感倾向", "积极情感倾向指数", "消极情感倾向指数"]
    ws.append(fields)
    for index, result in enumerate(results):
        print(results[index])
        #print('text: {},\t  predict: {}'.format(results[index]['text'][:50], results[index]['sentiment_key']))
        sentiment_key = results[index]['sentiment_key']
        if sentiment_key == "positive":
            positive_nums = positive_nums+1
            emotion = "积极"
        elif sentiment_key == "negative":
            negative_nums = negative_nums+1
            emotion = "消极"
        ws.append([results[index]['text'],  emotion, results[index]
                  ['positive_probs'], results[index]['negative_probs']])
    wb.save(filename=file_name)
    print('积极评论数: {},  消极评论数: {}'.format(positive_nums, negative_nums))
